from openai import OpenAI
from dotenv import load_dotenv
from app.utils.format_message import format_user_message
import tiktoken
import os
import aiohttp
import json
from typing import AsyncGenerator, Dict, Any, Optional
import time
import traceback

load_dotenv()


class OpenAIO1Service:
    def __init__(self):
        self.default_client = OpenAI(
            api_key=os.getenv("OPENAI_API_KEY"),
        )
        self.encoding = tiktoken.get_encoding("o200k_base")  # Encoder for OpenAI models
        self.base_url = "https://api.openai.com/v1/chat/completions"
        print(f"OpenAIO1Service 初始化完成，使用API密钥: {'已设置' if os.getenv('OPENAI_API_KEY') else '未设置'}")

    def call_o1_api(
        self,
        system_prompt: str,
        data: dict,
        api_key: str | None = None,
        model: str = "o1-mini",
    ) -> str:
        """
        Makes an API call to OpenAI o1-mini and returns the response.

        Args:
            system_prompt (str): The instruction/system prompt
            data (dict): Dictionary of variables to format into the user message
            api_key (str | None): Optional custom API key
            model (str): Model to use for the API call, defaults to "o1-mini"

        Returns:
            str: o1-mini's response text
        """
        # Create the user message with the data
        user_message = format_user_message(data)
        token_count = self.count_tokens(user_message) + self.count_tokens(system_prompt)
        print(f"[OpenAI o1] 非流式API调用: 总令牌数={token_count}, 使用模型={model}")

        # Use custom client if API key provided, otherwise use default
        client = OpenAI(api_key=api_key) if api_key else self.default_client

        try:
            print(
                f"[OpenAI o1] 使用{'自定义' if api_key else '默认'}API密钥调用{model}模型"
            )

            completion = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_message},
                ],
                max_completion_tokens=12000,  # Adjust as needed
                temperature=0.2,
            )

            response_text = completion.choices[0].message.content
            response_tokens = self.count_tokens(response_text)
            print(f"[OpenAI o1] API调用成功: 响应令牌数={response_tokens}")

            if response_text is None:
                raise ValueError("No content returned from OpenAI o1-mini")

            return response_text

        except Exception as e:
            print(f"[OpenAI o1] API调用错误: {str(e)}")
            raise

    async def call_o1_api_stream(
        self,
        system_prompt: str,
        data: dict,
        api_key: str | None = None,
        model: str = "o1-mini",
    ) -> AsyncGenerator[str, None]:
        """
        Makes a streaming API call to OpenAI o1-mini and yields the responses.

        Args:
            system_prompt (str): The instruction/system prompt
            data (dict): Dictionary of variables to format into the user message
            api_key (str | None): Optional custom API key
            model (str): Model to use for the API call, defaults to "o1-mini"

        Yields:
            str: Chunks of o1-mini's response text
        """
        # Create the user message with the data
        user_message = format_user_message(data)
        token_count = self.count_tokens(user_message) + self.count_tokens(system_prompt)
        print(f"[OpenAI o1] 流式API调用: 总令牌数={token_count}, 使用模型={model}")

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key or self.default_client.api_key}",
        }

        print(f"[OpenAI o1] 使用{'自定义' if api_key else '默认'}API密钥调用{model}流式模型")

        payload = {
            "model": model,
            "messages": [
                {
                    "role": "user",
                    "content": f"""
                    <VERY_IMPORTANT_SYSTEM_INSTRUCTIONS>
                    {system_prompt}
                    </VERY_IMPORTANT_SYSTEM_INSTRUCTIONS>
                    <USER_INSTRUCTIONS>
                    {user_message}
                    </USER_INSTRUCTIONS>
                    """,
                },
            ],
            "max_completion_tokens": 12000,
            "stream": True,
        }

        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    self.base_url, headers=headers, json=payload
                ) as response:

                    if response.status != 200:
                        error_text = await response.text()
                        print(f"Error response: {error_text}")
                        raise ValueError(
                            f"OpenAI API returned status code {response.status}: {error_text}"
                        )

                    line_count = 0
                    async for line in response.content:
                        line = line.decode("utf-8").strip()
                        if not line:
                            continue

                        line_count += 1

                        if line.startswith("data: "):
                            if line == "data: [DONE]":
                                break
                            try:
                                # 使用安全解析替代直接json.loads
                                json_str = line[6:]
                                data = self.safe_json_parse(json_str)
                                content = (
                                    data.get("choices", [{}])[0]
                                    .get("delta", {})
                                    .get("content")
                                )
                                if content:
                                    yield content
                            except Exception as e:
                                print(f"JSON处理错误: {e} for line: {line}")
                                continue

                    if line_count == 0:
                        print("[OpenAI o1] 警告: 流式响应中未收到任何行")
                    else:
                        print(f"[OpenAI o1] 流式API调用完成: 共接收 {line_count} 行数据")

        except aiohttp.ClientError as e:
            print(f"[OpenAI o1] 连接错误: {str(e)}")
            raise ValueError(f"Failed to connect to OpenAI API: {str(e)}")
        except Exception as e:
            print(f"[OpenAI o1] 流式API调用中出现意外错误: {str(e)}")
            raise

    def count_tokens(self, prompt: str) -> int:
        """
        Counts the number of tokens in a prompt.

        Args:
            prompt (str): The prompt to count tokens for

        Returns:
            int: Estimated number of input tokens
        """
        num_tokens = len(self.encoding.encode(prompt))
        return num_tokens

    # 添加一个安全JSON解析和序列化方法
    def safe_json_parse(self, data_str: str) -> Dict[str, Any]:
        """
        安全解析JSON字符串，处理可能的错误
        
        Args:
            data_str: 要解析的JSON字符串
            
        Returns:
            解析后的字典，如果解析失败则返回空字典
        """
        try:
            return json.loads(data_str)
        except json.JSONDecodeError as e:
            print(f"JSON解析错误: {e}, 尝试修复...")
            
            # 清理错误的Unicode字符
            cleaned_str = ''.join(c for c in data_str if ord(c) >= 32 and ord(c) < 127)
            
            try:
                return json.loads(cleaned_str)
            except json.JSONDecodeError:
                # 查找有效的JSON部分
                start = cleaned_str.find('{')
                end = cleaned_str.rfind('}')
                
                if start >= 0 and end > start:
                    try:
                        return json.loads(cleaned_str[start:end+1])
                    except:
                        pass
            
            # 所有恢复尝试失败
            print(f"无法解析JSON: {data_str[:100]}...")
            return {}
    
    def safe_json_dumps(self, data: Dict[str, Any]) -> str:
        """
        安全序列化字典为JSON字符串
        
        Args:
            data: 要序列化的字典
            
        Returns:
            序列化后的JSON字符串
        """
        try:
            # 预处理字典中的值，确保所有值都可以序列化
            def sanitize(obj):
                if isinstance(obj, dict):
                    return {k: sanitize(v) for k, v in obj.items()}
                elif isinstance(obj, list):
                    return [sanitize(v) for v in obj]
                elif isinstance(obj, (int, float, bool, str)) or obj is None:
                    return obj
                else:
                    # 将不可序列化的对象转换为字符串
                    return str(obj)
            
            sanitized_data = sanitize(data)
            return json.dumps(sanitized_data, ensure_ascii=False)
        except Exception as e:
            print(f"JSON序列化错误: {e}")
            # 失败时返回简化的JSON
            return json.dumps({"error": "数据序列化失败"})

    async def stream_openai_o1_mini(self, user_message: str, system_message: str = None, api_key: str = None):
        """
        流式调用Anthropic Claude模型API (o1-mini)
        
        Args:
            user_message: 用户消息
            system_message: 系统消息，默认为None
            api_key: API密钥，默认为None，使用默认API密钥
            
        Returns:
            异步生成器，生成Claude模型的响应流
        """
        try:
            # 计算标记数
            token_count = self.count_tokens(user_message)
            if system_message:
                token_count += self.count_tokens(system_message)
            
            print(f"[Claude] 流式API调用: 总令牌数={token_count}")
            print(f"[Claude] 使用{'自定义' if api_key else '默认'}API密钥调用Claude流式模型")
            
            # 获取API密钥
            api_key = api_key if api_key else self.default_client.api_key
            
            # 准备消息
            messages = []
            if system_message:
                messages.append({"role": "system", "content": system_message})
            messages.append({"role": "user", "content": user_message})
            
            # 记录开始时间
            start_time = time.time()
            
            # Anthropic API调用参数  
            data = {
                "model": "claude-3-opus-20240229",
                "messages": messages,
                "stream": True
            }
            
            # API调用
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    "https://api.anthropic.com/v1/messages",
                    headers={
                        "x-api-key": api_key,
                        "anthropic-version": "2023-06-01",
                        "content-type": "application/json"
                    },
                    json=data
                ) as response:
                    # 检查响应状态
                    response.raise_for_status()
                    
                    line_count = 0
                    content_buffer = ""
                    
                    # 处理SSE流式响应
                    async for line in response.content:
                        line_count += 1
                        line = line.decode('utf-8').strip()
                        if line and line != "":
                            if line.startswith("data: "):
                                line = line[6:]  # 移除 "data: " 前缀
                                if line == "[DONE]":
                                    break
                                
                                try:
                                    # 使用安全解析替代直接json.loads
                                    j = self.safe_json_parse(line)
                                    if j.get("type") == "content_block_delta":
                                        if "delta" in j and "text" in j.get("delta", {}):
                                            content = j["delta"]["text"]
                                            content_buffer += content
                                            yield content
                                except Exception as e:
                                    print(f"[Claude] 警告: 无法解析JSON: {line}, 错误: {e}")
                                    continue
                    
                    # 记录结束时间和总耗时
                    end_time = time.time()
                    elapsed_time = end_time - start_time
                    
                    if line_count == 0:
                        print("[Claude] 警告: 流式响应中未收到任何行")
                        
                    # 计算返回的内容长度的标记数    
                    response_token_count = self.count_tokens(content_buffer)
                    
                    print(f"[Claude] 流式API调用完成: 共接收 {line_count} 行数据")
                    print(f"[Claude] 响应标记数: {response_token_count}, 总耗时: {elapsed_time:.2f}秒")
                    
        except aiohttp.ClientError as e:
            print(f"[Claude] 连接错误: {str(e)}")
            yield f"API连接错误: {str(e)}"
        except Exception as e:
            print(f"[Claude] 流式API调用中出现意外错误: {str(e)}")
            traceback.print_exc()
            yield f"API调用错误: {str(e)}"
